#! /usr/bin/env python3
# -*-coding: utf-8-*-


__author__ = "Moonkie"


import os
import sys
import glob
import numpy as np
import cv2

# from svd import svd

# img_match = '*[.bmp|.png|.jpg|.jpeg]'
CURDIR = os.path.split(os.path.abspath(__file__))[0]
ROOTDIR = os.path.split(CURDIR)[0]
DATASETS = os.path.join(ROOTDIR,'datasets')
TRAINDB = os.path.join(DATASETS,'trainned.npz')
TRAINSRC = os.path.join(DATASETS,'svds')
IMGMATCH = '*[.bmp|.png|.jpg|.pgm]'
IMAGEFILTER = ('jpg','bmp','jpg','pgm','png')

sys.path.append(os.path.join(ROOTDIR,'db'))
from db import Chip,People

class KlError(Exception):
    def __init__(self,msg,point):
        self.msg = "[QAQ] %s.(%s in svd/kl.py)" % (msg,point)
    def __str__(self):
        return self.msg



class KL(object):
    '''KL pca

    '''
    def __init__(self,train_folder=TRAINSRC,limit_folder=True,opt_folder=None):
        # self.train_images = self.load_train_image(train_folder=train_folder,limit_folder=limit_folder,opt_folder=opt_folder)
        self.to_update_config = (train_folder,limit_folder,opt_folder)
    
    @staticmethod
    def trainimg2vects(trainningvects):
        '''Format trainning images to vector

        transform image these will be trainned
        '''
        S = []; shape=[]; filepath = dict()
        if bool(trainningvects) and len(trainningvects) > 0:
            for i in range(len(trainningvects)):
                image = trainningvects[i]
                filepath[i] = image
                img = cv2.imread(image)
                if img.ndim == 2:
                    rows,cols = img.shape
                    shape.append((rows,cols))
                    img = np.reshape(img,(rows*cols))
                elif img.ndim == 3:
                    rows,cols,dims = img.shape
                    shape.append((rows,cols,dims))
                    img = np.reshape(img,(rows*cols*dims))
                else:
                    raise KlError('This image format not be supported now','trainimg2vects') #Exception('[QAQ] .(img2vects in svd/kl.py)')
                S.append(list(img))
            S = np.array(S)
            # S = np.vstack(S)
            return S,shape,filepath
        else:
            return np.array(S)

    def testimg2vect(self,imagepath):
        img = cv2.imread(imagepath)
        if img.ndim == 2:
            rows,cols = img.shape
            img = np.reshape(img,(rows*cols))
        elif img.ndim == 3:
            rows,cols,dims = img.shape
            img = np.reshape(img,(rows*cols*dims))
        else:
            raise KlError('Not support this image pattern','testimg2vect')
        return img
    
    @staticmethod
    def load_train_image(train_folder,limit_folder,opt_folder):
        '''Load tranning images from stroage devices

        '''
        train_images = []
        if bool(limit_folder):
            train_folder = os.path.join(train_folder,IMGMATCH)
            train_images.extend(glob.glob(train_folder))
        elif not opt_folder or len(opt_folder) == 0: 
            for root,dirs,files in os.walk(train_folder,topdown=True):
                if files and len(files) > 0:
                    for file in files:
                        if os.path.split(file)[-1].lower().endswith(IMAGEFILTER):
                            train_images.append(os.path.join(root,file))
        elif opt_folder and len(opt_folder) != 0:
            # print("test")
            for d in opt_folder:
                folderpath = os.path.join(train_folder,d,IMGMATCH)
                train_images.extend(glob.glob(folderpath))
        return train_images
    
    def train(self):
        '''trainning images method

        '''
        if not hasattr(self,'imagevects'):
            raise KlError('Images not given or not update(exec "update" before "train") or load trained database?(use "load" and not need to "train" any more)','train')
        if self.imagevects.any() and self.imagevects.itemsize > 0:
            self.traindata,self.pretest = kl_transfrom(self.imagevects)
        else:
            raise KlError('Not have images to train','train')

    def verify(self,testimgpath):
        '''verify image from train data

        this method base on Euclidean distance.
        '''
        if not self.traindata.any() and self.traindata.itemsize == 0:
            raise KlError('Not found any trainned image data','verify')
        if not os.path.isfile(testimgpath)  or not str.endswith(testimgpath.lower(),IMAGEFILTER):
            raise KlError('Cant operate this file %s' % testimgpath,'verify')
        timg = self.testimg2vect(testimgpath)
        timg = np.dot(timg,self.pretest)
        rsl = np.argmin([np.linalg.norm(i - timg) for i in self.traindata.T])
        image = self.imgfilepath[rsl] 
        return image

    def save(self,db=TRAINDB):
        if not hasattr(self,'traindata'):
            raise KlError("Nothing need to save.Must be train before",'save')
        np.savez_compressed(db,traindata=self.traindata,pretest=self.pretest,imgfilepath=self.imgfilepath)
        print('[^_^] Sava trainned data successfully.')

    def load(self,db=TRAINDB):
        _ = np.load(db)
        self.traindata = _["traindata"]
        self.pretest = _["pretest"]
        self.imgfilepath = _["imgfilepath"].tolist()
        # print(self.imgfilepath)

    def update(self):
        if not bool(self.to_update_config) and len(self.to_update_config) != 3:
            raise KlError('Bad updating configurations are given','update')
        trainningvects = KL.load_train_image(self.to_update_config[0],self.to_update_config[1],self.to_update_config[2])
        self.imagevects,self.imageshape,self.imgfilepath = KL.trainimg2vects(trainningvects)
    
    @staticmethod
    def pca_compress_imgfile(src,dest,rate=.8,folder_limited=True,opt_folder=None):
        from pca_compress import pca_compress
        imagelist = KL.load_train_image(src,folder_limited,opt_folder)
        imgvects,imgshape,imgfilepath = KL.trainimg2vects(imagelist)
        print(imgvects)
        imgvects = np.float64(imgvects)
        nldata,rdvects,mean = pca_compress(imgvects.T,rate)
        np.savez_compressed(dest,nldata=nldata,rdvects=rdvects,mean=mean,imgshape=imgshape)
        print('[^_^] Compress successfully.')
                

def kl_transfrom(data):
    '''K-L 变换实现

    基于主成分分析做K-L变化
    '''
    # print('test','\n',data.shape)
    if not data.any() or data.itemsize == 0:
        raise KlError('May be you give some fatal data','kl_transform')
    data = np.float64(data)
    return pca(data)


def pca(data):
    '''PCA 算法特征提取和模式识别

    '''
    symvects = np.dot(data,data.T)  
    eigvals,eigvects = np.linalg.eig(symvects)  
    matvects = np.dot(data.T,eigvects)  
    trainvects = np.dot(matvects.T,data.T)
    return trainvects,matvects


if __name__ == '__main__':
    tests = os.path.join(DATASETS,'faces')
    # kl = KL('1')
    # kl.update()   # 使KL初始化参数生效
    # kl.train()    # 开始训练

    # kl.save()     # 保存训练数据 (可不使用)

    # kl.load()     # 加载训练数据库 (加载数据库则 update方法可以不调用，也不需要训练模型)

    # image = kl.verify('test.BMP')     # 校验
    # print(image)

    # KL.pca_compress_imgfile('C:\\Users\\moonkie\\Desktop\\test','C:\\Users\\moonkie\\Desktop\\compress.npz')
    # 默认不加参数训练g
    kl = KL()
    kl.update()
    kl.train()
    # kl.save()
    # kl.load()
    # chip = kl.verify(os.path.join(tests,'one.png'))
    # print(chip) 
    pass